from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score, classification_report, roc_curve, auc
import label_encode
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
import joblib
from xgboost import XGBClassifier

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号 '-' 显示为方块的问题

# 读取数据
data = pd.read_csv("../../data/raw/train.csv")
# 数据预处理  将有多分类的列 转化为 数值型的
# 需要编码的列
cols = [
    "BusinessTravel", "Department", "Education", "EducationField",
    "EnvironmentSatisfaction", "Gender", "JobInvolvement", "JobLevel", "JobRole",
    "JobSatisfaction", "MaritalStatus", "OverTime", "PerformanceRating",
    "RelationshipSatisfaction", "TrainingTimesLastYear", "StockOptionLevel", "WorkLifeBalance"
]
data = label_encode.encode(data, cols)
# 选取特征列  12 + 9列
# 上边通过卡方检验 选出了 12列
# 下边通过 随机森林 + 置换重要性选出了 9列
X = data[[
    "BusinessTravel",
    "Department",
    "EducationField",
    "EnvironmentSatisfaction",
    "JobInvolvement",
    "JobLevel",
    "JobRole",
    "JobSatisfaction",
    "MaritalStatus",
    "OverTime",
    "StockOptionLevel",
    "WorkLifeBalance",

    "DistanceFromHome",
    "Age",
    "YearsInCurrentRole",
    "YearsSinceLastPromotion",
    "MonthlyIncome",
    "PercentSalaryHike",
    "YearsWithCurrManager",
    "NumCompaniesWorked",
    "TotalWorkingYears",
]
]
y = data["Attrition"]
# 对X 标准化
scale = StandardScaler()
scale.fit(X)
X = scale.transform(X)
joblib.dump(scale, "scale_model.pkl")

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=25)
# 使用XGBClassifier 进行模型训练
es_xgb = XGBClassifier(
    n_estimators=500,
    max_depth=1,  # 2
    min_child_weight=3,
    reg_lambda=1.0,
    objective='binary:logistic',
    early_stopping_rounds=50,
    eval_metric='auc',
)

es_xgb.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_pred = es_xgb.predict_proba(X_test)[:, 1]
# y_pred_1 = es_xgb.predict(X_test)   # 用于分类评估报告

fpr, tpr, thresholds = roc_curve(y_test, y_pred)
roc_auc = auc(fpr, tpr)  # 或直接用roc_auc_score

# 4. 绘制ROC曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')  # 随机线
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic (ROC) Curve')
plt.legend(loc="lower right")
plt.grid(True)
plt.savefig("xgb_roc_curve.png")

plt.show()

# 打印分类评估报告
# print("分类评估报告：")
# print(classification_report(y_test, y_pred_1))
print("AUC:", roc_auc_score(y_test, y_pred))  # AUC: 0.9092163555813785

# 保存模型
joblib.dump(es_xgb, "xgb_model.pkl")
